import os
import cv2
import dlib
import shutil
import re
import numpy as np

# 设置源目录和目标目录
source_dir = 'source'  # 输入文件夹
out_dir = 'out'  # 输出文件夹
to_dir = os.path.join(out_dir, '已遮脸')  # 检测到人脸的输出文件夹
noto_dir = os.path.join(out_dir, '未遮脸')  # 未检测到人脸的输出文件夹

# 确保目标文件夹存在
for directory in [to_dir, noto_dir]:
    if not os.path.exists(directory):
        os.makedirs(directory)

# 加载 dlib 的人脸检测器
detector = dlib.get_frontal_face_detector()

# 加载遮挡用的PNG图片（带透明背景的PNG）
overlay = cv2.imread('face_mask.png', cv2.IMREAD_UNCHANGED)

 
# 函数：读取支持中文路径的图片
def imread_with_chinese_path(path):
    # 将文件路径转换为二进制数据
    img_data = np.fromfile(path, dtype=np.uint8)
    # 用 cv2.imdecode 来读取图像数据
    img = cv2.imdecode(img_data, cv2.IMREAD_UNCHANGED)
    return img

# 递归重命名源目录中的文件
 

# 遍历源目录中的所有文件
for filename in os.listdir(source_dir):
    # 只处理图片文件（根据扩展名过滤）
    if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
        img_path = os.path.join(source_dir, filename)

        # 读取原图像，使用自定义函数以支持中文路径
        img = imread_with_chinese_path(img_path)

        # 检查图像是否成功加载
        if img is None:
            print(f"无法加载图像: {img_path}")
            continue  # 跳过此文件并继续处理下一个文件

        # 将图像转换为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 检测人脸
        faces = detector(gray)

        # 检测到人脸的情况
        if len(faces) > 0:
            # 遍历检测到的人脸
            for face in faces:
                # 增加一定的边距，使得遮挡区域稍微大于检测到的人脸
                x, y, w, h = face.left(), face.top(), face.width(), face.height()
                margin = 50  # 边距大小，适当增大
                x = max(0, x - margin)  # 防止超出边界
                y = max(0, y - margin)

                # 向上移动的偏移量
                offset = 15  # 向上移动的像素数
                y = max(0, y - offset)  # 确保不超出边界

                w += margin * 2  # 增加宽度
                h += margin * 2  # 增加高度

                # 调整遮挡图片的大小以匹配脸部大小
                resized_overlay = cv2.resize(overlay, (w, h))

                # 分离PNG图片的BGR与Alpha通道（透明度）
                overlay_rgb = resized_overlay[:, :, :3]  # BGR通道
                overlay_alpha = resized_overlay[:, :, 3]  # Alpha通道

                # 获取脸部区域的ROI（区域部分）
                roi = img[y:y+h, x:x+w]

                # 确保 ROI 的大小与遮挡图像相匹配
                roi_resized = cv2.resize(roi, (w, h))

                # 根据 alpha 通道创建遮罩，保留透明部分
                mask = overlay_alpha / 255.0  # 将alpha通道归一化到0-1之间
                inv_mask = 1.0 - mask  # 反转遮罩，用于背景

                # 用 alpha 通道合成图片：前景(overlay)和背景(roi_resized)叠加
                for c in range(0, 3):  # 对 BGR 三个通道操作
                    roi_resized[:, :, c] = (mask * overlay_rgb[:, :, c] + inv_mask * roi_resized[:, :, c])

                # 确保在合成前检查形状
                if roi_resized.shape == img[y:y+h, x:x+w].shape:
                    img[y:y+h, x:x+w] = roi_resized
                else:
                    print(f"Shape mismatch for {filename}: ROI {roi_resized.shape} does not match target {img[y:y+h, x:x+w].shape}.")

            # 保存处理后的图像到已遮脸目录，保持质量
            output_path = os.path.join(to_dir, filename)
            cv2.imencode('.jpg', img)[1].tofile(output_path)  # 使用 .tofile 保存中文路径的文件
            print(f"检测到人脸，已保存到: {output_path}")

        else:
            # 直接将未检测到人脸的图像复制到未遮脸目录
            noface_path = os.path.join(noto_dir, filename)
            shutil.copy(img_path, noface_path)
            print(f"未检测到人脸，已复制到: {noface_path}")

print("处理完成！所有图像已保存到目标目录。")
